// Tencent is pleased to support the open source community by making
// 蓝鲸智云 - 监控平台/日志平台 (BlueKing - Monitor/Log) available.
// Copyright (C) 2017-2022 THL A29 Limited, a Tencent company. All rights reserved.
// Licensed under the MIT License (the "License"); you may not use this file except in compliance with the License.
// You may obtain a copy of the License at http://opensource.org/licenses/MIT
// Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.
//

package resourcefilter

import (
	"strings"

	"github.com/TencentBlueKing/bkmonitor-kits/logger"
	"github.com/mitchellh/mapstructure"
	"go.opentelemetry.io/collector/model/pdata"

	"github.com/TencentBlueKing/bk-collector/confengine"
	"github.com/TencentBlueKing/bk-collector/define"
	"github.com/TencentBlueKing/bk-collector/processor"
)

const (
	ProcessorName string = "resource_filter"
)

func init() {
	processor.Register(ProcessorName, NewFactory)
}

func NewFactory(conf map[string]interface{}, customized []processor.SubConfigProcessor) (processor.Processor, error) {
	return newFactory(conf, customized)
}

func newFactory(conf map[string]interface{}, customized []processor.SubConfigProcessor) (*resourceFilter, error) {
	configs := confengine.NewTierConfig()

	c := &Config{}
	if err := mapstructure.Decode(conf, c); err != nil {
		return nil, err
	}
	c.Clean()
	configs.SetGlobal(*c)

	for _, custom := range customized {
		cfg := &Config{}
		if err := mapstructure.Decode(custom.Config.Config, cfg); err != nil {
			logger.Errorf("failed to decode config: %v", err)
			continue
		}
		cfg.Clean()
		configs.Set(custom.Token, custom.Type, custom.ID, *cfg)
	}

	return &resourceFilter{
		CommonProcessor: processor.NewCommonProcessor(conf, customized),
		configs:         configs,
	}, nil
}

type resourceFilter struct {
	processor.CommonProcessor
	configs *confengine.TierConfig // type: Config
}

func (p resourceFilter) Name() string    { return ProcessorName }
func (p resourceFilter) IsDerived() bool { return false }

func (p resourceFilter) Process(record *define.Record) (*define.Record, error) {
	config := p.configs.GetByToken(record.Token.Original).(Config)
	if len(config.Replace) > 0 {
		p.replaceAction(record)
	}
	if len(config.Add) > 0 {
		p.addAction(record)
	}
	if len(config.Assemble) > 0 {
		p.assembleAction(record)
	}
	if len(config.Drop.Keys) > 0 {
		p.dropAction(record)
	}
	return nil, nil
}

func (p resourceFilter) assembleAction(record *define.Record) {
	switch record.RecordType {
	case define.RecordTraces:
		pdTraces := record.Data.(pdata.Traces)
		resourceSpans := pdTraces.ResourceSpans()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Assemble {
			for i := 0; i < resourceSpans.Len(); i++ {
				spans := resourceSpans.At(i)
				attributes := spans.Resource().Attributes()
				var values []string
				for _, key := range action.Keys {
					v, ok := attributes.Get(key)
					if !ok {
						// 空值保留
						values = append(values, "")
						continue
					}
					values = append(values, v.AsString())
				}
				attributes.UpsertString(action.Destination, strings.Join(values, action.Separator))
			}
		}
	}
}

func (p resourceFilter) addAction(record *define.Record) {
	switch record.RecordType {
	case define.RecordTraces:
		pdTraces := record.Data.(pdata.Traces)
		resourceTraces := pdTraces.ResourceSpans()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Add {
			for i := 0; i < resourceTraces.Len(); i++ {
				trace := resourceTraces.At(i)
				trace.Resource().Attributes().UpsertString(action.Label, action.Value)
			}
		}

	case define.RecordMetrics:
		pdMetrics := record.Data.(pdata.Metrics)
		resourceMetrics := pdMetrics.ResourceMetrics()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Add {
			for i := 0; i < resourceMetrics.Len(); i++ {
				metric := resourceMetrics.At(i)
				metric.Resource().Attributes().UpsertString(action.Label, action.Value)
			}
		}

	case define.RecordLogs:
		pdLogs := record.Data.(pdata.Logs)
		resourceLogs := pdLogs.ResourceLogs()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Add {
			for i := 0; i < resourceLogs.Len(); i++ {
				log := resourceLogs.At(i)
				log.Resource().Attributes().UpsertString(action.Label, action.Value)
			}
		}
	}
}

func (p resourceFilter) dropAction(record *define.Record) {
	switch record.RecordType {
	case define.RecordTraces:
		pdTraces := record.Data.(pdata.Traces)
		resourceTraces := pdTraces.ResourceSpans()
		for _, dimension := range p.configs.GetByToken(record.Token.Original).(Config).Drop.Keys {
			for i := 0; i < resourceTraces.Len(); i++ {
				trace := resourceTraces.At(i)
				trace.Resource().Attributes().Delete(dimension)
			}
		}

	case define.RecordMetrics:
		pdMetrics := record.Data.(pdata.Metrics)
		resourceMetrics := pdMetrics.ResourceMetrics()
		for _, dimension := range p.configs.GetByToken(record.Token.Original).(Config).Drop.Keys {
			for i := 0; i < resourceMetrics.Len(); i++ {
				metric := resourceMetrics.At(i)
				metric.Resource().Attributes().Delete(dimension)
			}
		}

	case define.RecordLogs:
		pdLogs := record.Data.(pdata.Logs)
		resourceLogs := pdLogs.ResourceLogs()
		for _, dimension := range p.configs.GetByToken(record.Token.Original).(Config).Drop.Keys {
			for i := 0; i < resourceLogs.Len(); i++ {
				log := resourceLogs.At(i)
				log.Resource().Attributes().Delete(dimension)
			}
		}
	}
}

func (p resourceFilter) replaceAction(record *define.Record) {
	switch record.RecordType {
	case define.RecordTraces:
		pdTraces := record.Data.(pdata.Traces)
		resourceTraces := pdTraces.ResourceSpans()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Replace {
			for i := 0; i < resourceTraces.Len(); i++ {
				trace := resourceTraces.At(i)
				v, ok := trace.Resource().Attributes().Get(action.Source)
				if !ok {
					continue
				}
				trace.Resource().Attributes().Delete(action.Source)
				trace.Resource().Attributes().Upsert(action.Destination, v)
			}
		}

	case define.RecordMetrics:
		pdMetrics := record.Data.(pdata.Metrics)
		resourceMetrics := pdMetrics.ResourceMetrics()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Replace {
			for i := 0; i < resourceMetrics.Len(); i++ {
				metric := resourceMetrics.At(i)
				v, ok := metric.Resource().Attributes().Get(action.Source)
				if !ok {
					continue
				}
				cloned := pdata.NewAttributeValueEmpty()
				v.CopyTo(cloned)
				metric.Resource().Attributes().Delete(action.Source)
				metric.Resource().Attributes().Upsert(action.Destination, cloned)
			}
		}

	case define.RecordLogs:
		pdLogs := record.Data.(pdata.Logs)
		resourceLogs := pdLogs.ResourceLogs()
		for _, action := range p.configs.GetByToken(record.Token.Original).(Config).Replace {
			for i := 0; i < resourceLogs.Len(); i++ {
				log := resourceLogs.At(i)
				v, ok := log.Resource().Attributes().Get(action.Source)
				if !ok {
					continue
				}
				log.Resource().Attributes().Delete(action.Source)
				log.Resource().Attributes().Upsert(action.Destination, v)
			}
		}
	}
}
